import numpy as np
import cmath 
import math 
import matplotlib.pyplot as plt
from scipy.fft import fft, fftfreq


def generate_cos_wave(freq, sample_rate, duration):
    x = np.linspace(0, duration, sample_rate * duration, endpoint=False)
    frequencies = x * freq
    # 2pi because np.sin takes radians
    # 2*pi*f*n*T
    y = np.cos((2 * np.pi) * frequencies)
    return x, y


def generate_rotation_factor(N, k):
	e_k = np.arange(N, dtype = complex)
	for n in range(N):
		e_k[n] = cmath.exp(complex(real=0, imag=-1)*(2*cmath.pi/N)*k*n)
	return e_k

def dft_transform(xn, N):
	#pass
	Xk = np.arange(N, dtype = complex)
	for n in range(N):
		Xk[n] = generate_rotation_factor(N, n).dot(xn)
	return Xk


def dft_freq(N, d):
	fk_negative = np.arange(start=-N/2,stop=0,step=1)*(1/(N*d))
	fk_positive = np.arange(start=0,stop=N/2,step=1)*(1/(N*d))
	fk = np.concatenate((fk_positive, fk_negative))
	return fk 


def main():
	sample_rate = 130
	duration = 3
	x, xn = generate_cos_wave(2, sample_rate, duration)
	Xk = dft_transform(sample_rate*duration, xn)
	fk = dft_freq(sample_rate*duration, 1/sample_rate)
	print(Xk)
	plt.stem(fk, abs(Xk))
	plt.show()

	# plt.stem(xn)
	# plt.show()

	Xk = fft(xn)
	print("----------------------------")
	print(Xk)
	xf = fftfreq(sample_rate * duration, 1 / sample_rate)

	#print(Xk)
	print(xf)

	plt.stem(xf, abs(Xk))
	plt.show()

	# x = complex(0, 1)
	# x_e = cmath.exp(x)
	# xn = []
	# print(x,x_e)
	# print("cos(1)={},sin(1)={}".format(math.cos(1), math.sin(1)))


#main()
# generate_rotation_factor(10, 1)
# a = np.array([1,2,3])
# b = np.array([1,2,3])
# print((a*b).sum())
# print(a.dot(b))